import websocket
import json
import ssl
import time
import hashlib
import base64
import hmac
from datetime import datetime
from time import mktime
from urllib.parse import urlencode
from wsgiref.handlers import format_date_time


def test_websocket_connection():
    """测试WebSocket连接"""

    # 您的API配置
    app_id = "0adb1398"
    api_secret = "UJnvocAcTLWdATUCivPv:YnsopnQsfkJanwuLoFCV"
    api_key = "4e0ab6929faad9ff81b6b84d0fc8bf76"

    host = "spark-api.xf-yun.com"
    path = "/v4.0/chat"

    def create_signature():
        """创建签名"""
        now = datetime.now()
        date = format_date_time(mktime(now.timetuple()))

        signature_origin = f"host: {host}\ndate: {date}\nGET {path} HTTP/1.1"
        signature_sha = hmac.new(
            api_secret.encode('utf-8'),
            signature_origin.encode('utf-8'),
            digestmod=hashlib.sha256
        ).digest()
        signature_sha_base64 = base64.b64encode(signature_sha).decode('utf-8')

        authorization_origin = f'api_key="{api_key}", algorithm="hmac-sha256", headers="host date request-line", signature="{signature_sha_base64}"'
        authorization = base64.b64encode(authorization_origin.encode('utf-8')).decode('utf-8')

        return authorization, date

    authorization, date = create_signature()

    # 构建URL
    params = {
        "authorization": authorization,
        "date": date,
        "host": host
    }
    query_string = urlencode(params)
    url = f"wss://{host}{path}?{query_string}"

    print(f"连接URL: {url}")

    received_messages = []

    def on_message(ws, message):
        print(f"收到消息: {message}")
        received_messages.append(message)

        try:
            data = json.loads(message)
            code = data.get('header', {}).get('code', -1)
            if code != 0:
                print(f"错误代码: {code}, 消息: {data.get('header', {}).get('message', '未知错误')}")
            else:
                choices = data.get("payload", {}).get("choices", {})
                status = choices.get("status", -1)
                content = choices.get("text", [{}])[0].get("content", "")
                print(f"AI回复: {content}")

                if status == 2:
                    print("对话完成")
                    ws.close()
        except Exception as e:
            print(f"解析消息错误: {e}")

    def on_error(ws, error):
        print(f"WebSocket错误: {error}")

    def on_close(ws, close_status_code, close_msg):
        print("连接关闭")

    def on_open(ws):
        print("连接已建立")
        data = {
            "header": {
                "app_id": app_id,
                "uid": "test_user"
            },
            "parameter": {
                "chat": {
                    "domain": "4.0Ultra",
                    "temperature": 0.5,
                    "max_tokens": 2048
                }
            },
            "payload": {
                "message": {
                    "text": [
                        {"role": "user", "content": "你好，请简单介绍一下你自己"}
                    ]
                }
            }
        }
        ws.send(json.dumps(data))
        print("已发送测试消息")

    try:
        ws = websocket.WebSocketApp(
            url,
            on_message=on_message,
            on_error=on_error,
            on_close=on_close,
            on_open=on_open
        )

        print("开始连接...")
        ws.run_forever(sslopt={"cert_reqs": ssl.CERT_NONE})

        return len(received_messages) > 0

    except Exception as e:
        print(f"连接异常: {e}")
        return False


if __name__ == "__main__":
    success = test_websocket_connection()
    if success:
        print("API连接测试成功！")
    else:
        print("API连接测试失败，请检查：")
        print("1. API密钥是否正确")
        print("2. 网络连接是否正常")
        print("3. 账户是否已开通服务")
        print("4. 账户余额是否充足")